﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Text;

namespace dnsServerCSharp
{

    // 查询码
    public enum DnsQueryType
    {
        Query = 0,              // Query
        IQuery = 1,             // Obsolete: IP to name
        ServerStatus = 2,       // Obsolete: DNS ping
        Unknown = 3,            // Unknown
        Notify = 4,             // Notify
        Update = 5              // Dynamic Update
    }

    // 响应码（错误代号）
    public enum DnsResponseCode
    {
        NoError = 0,
        FormatError = 1,        // Format error
        ServerFailure = 2,      // Server failure
        NameError = 3,          // Name error
        NotImplemented = 4,     // Not implemented
        Refused = 5,            // Refused
        DomainNotExist = 6,     // Domain name should not exist
        RRShouldNotExist = 7,   // RR set should not exist
        RRNetExist = 8,         // RR set does not exist
        NotAuthoritative = 9,   // Not authoritative for zone
        NotZone = 10,           // Name is not zone
    }
    public enum DnsClass
    {
        Internet = 0x0001,
        Chaos = 0x0003,
        Hesiod = 0x0004,
        None = 0x00FE,
        All = 0x00FF,
        Any = 0x00FF
    }
    public enum DnsRecordType
    {
        A = 0x0001,
        NS = 0x0002,
        MD = 0x0003,
        MF = 0x0004,
        CNAME = 0x0005,
        SOA = 0x0006,
        MB = 0x0007,
        MG = 0x0008,
        MR = 0x0009,
        NULL = 0x000a,
        WKS = 0x000b,
        PTR = 0x000c,
        HINFO = 0x000d,
        MINFO = 0x000e,
        MX = 0x000f,
        TEXT = 0x0010,
        AAAA = 0x001C,
        A6 = 0x0026,
    }
    public class DnsHeader
    {
        public ushort id = 0;

        // operation byte
        public bool isResponse = false;
        public DnsQueryType opcode = DnsQueryType.Query;
        public bool authoritative = false;
        public bool truncation = false;
        public bool recursionDesired = true;

        // response byte
        public bool recursionAvailable = true;
        public bool reserved = false;
        public bool authenticatedData = false;
        public bool checkingDisabled = false;
        public DnsResponseCode responseCode = DnsResponseCode.NoError;

        //ushort questionCount;
        //ushort answerCount;
        //ushort nameServerCount;
        //ushort additionalCount;

        public byte operationByte
        {
            get
            {
                return (byte)(
                    Convert.ToByte(isResponse) << 7 |
                    ((byte)opcode & 0x0F) << 3 |
                    Convert.ToByte(authoritative) << 2 |
                    Convert.ToByte(truncation) << 1 |
                    Convert.ToByte(recursionDesired)
                    );
            }
            set
            {
                isResponse = Convert.ToBoolean(value & (1 << 7));
                opcode = (DnsQueryType)((value >> 3) & 0x0F);
                authoritative = Convert.ToBoolean(value & (1 << 2));
                truncation = Convert.ToBoolean(value & (1 << 1));
                recursionDesired = Convert.ToBoolean(value & 1);
            }
        }

        public byte responseByte
        {
            get
            {
                return (byte)(
                    Convert.ToByte(recursionAvailable) << 7 |
                    Convert.ToByte(reserved) << 6 |
                    Convert.ToByte(authenticatedData) << 5 |
                    Convert.ToByte(checkingDisabled) << 4 |
                    ((byte)responseCode & 0x0F)
                    );
            }
            set
            {
                recursionAvailable = Convert.ToBoolean(value & (1 << 7));
                reserved = Convert.ToBoolean(value & (1 << 6));
                authenticatedData = Convert.ToBoolean(value & (1 << 5));
                checkingDisabled = Convert.ToBoolean(value & (1 << 4));
                responseCode = (DnsResponseCode)(value & 0x0F);
            }
        }

        //! deep copy
        public DnsHeader Clone()
        {
            var obj = new DnsHeader();
            obj.id = this.id;
            obj.operationByte = this.operationByte;
            obj.responseByte = this.responseByte;
            return obj;
        }
    }

    public class DnsQuestion
    {
        public string name { set; get; }
        public DnsRecordType recordType { set; get; }
        public DnsClass dnsClass { set; get; }

        //! deep copy
        public DnsQuestion Clone()
        {
            var obj = new DnsQuestion();
            obj.name = this.name;
            obj.recordType = this.recordType;
            obj.dnsClass = this.dnsClass;
            return obj;
        }
    }
    public class DnsResource
    {
        public string name { set; get; }
        public DnsRecordType recordType { set; get; }
        public DnsClass dnsClass { set; get; }
        public uint timeToLive { set; get; }
        public ushort dataLength { set; get; }
        public byte[] data { set; get; }

        public void setData(string name, DnsRecordType recordType, DnsClass dnsClass, uint timeToLive, byte[] data)
        {
            this.name = name;
            this.recordType = recordType;
            this.dnsClass = dnsClass;
            this.timeToLive = timeToLive;
            this.data = data;
            this.dataLength = (ushort)data.Length;
        }

        //! deep copy
        public DnsResource Clone()
        {
            var obj = new DnsResource();
            obj.name = this.name;
            obj.recordType = this.recordType;
            obj.dnsClass = this.dnsClass;
            obj.timeToLive = this.timeToLive;
            obj.dataLength = this.dataLength;
            obj.data = new byte[this.dataLength];
            this.data.CopyTo(obj.data, 0);
            return obj;
        }

        // for debug
        DnsAData aData
        {
            get
            {
                return DnsAData.fromData(data);
            }
        }
        DnsAAAAData aaaaData
        {
            get
            {
                return DnsAAAAData.fromData(data);
            }
        }
        DnsNsData nsData
        {
            get
            {
                return DnsNsData.fromData(new byte[0], data);
            }
        }
        DnsCNameData cnameData
        {
            get
            {
                return DnsCNameData.fromData(new byte[0], data);
            }
        }
    }
    public class DnsAData
    {
        public System.Net.IPAddress address { set; get; }

        public static DnsAData fromData(byte[] data)
        {
            DnsAData obj = new DnsAData();
            obj.address = new System.Net.IPAddress(data);
            return obj;
        }

        public byte[] toData()
        {
            return address.GetAddressBytes();
        }
    }
    public class DnsAAAAData
    {
        public System.Net.IPAddress address { set; get; }

        public static DnsAAAAData fromData(byte[] data)
        {
            var newBytes = new byte[8];
            Array.Copy(data, 0, newBytes, 4, 4);
            Array.Copy(data, 4, newBytes, 0, 4);

            DnsAAAAData obj = new DnsAAAAData();
            obj.address = new System.Net.IPAddress(newBytes);
            return obj;
        }

        public byte[] toData()
        {
            var bytes = address.GetAddressBytes();
            var retv = new byte[8];
            Array.Copy(bytes, 0, retv, 4, 4);
            Array.Copy(bytes, 4, retv, 0, 4);
            return retv;
        }
    }

    public class DnsNsData
    {
        public string name { set; get; }
        public static DnsNsData fromData(byte[] dataHead, byte[] data)
        {
            DnsNsData obj = new DnsNsData();
            obj.name = DnsMessage.parseDnsStandardName(dataHead, data);
            return obj;
        }
        public byte[] toData()
        {
            return DnsMessage.stringToDnsStandardName(name);
        }
    }

    public class DnsCNameData
    {
        public string name { set; get; }
        public static DnsCNameData fromData(byte[] dataHead, byte[] data)
        {
            DnsCNameData obj = new DnsCNameData();
            obj.name = DnsMessage.parseDnsStandardName(dataHead, data);
            return obj;
        }
        public byte[] toData()
        {
            return DnsMessage.stringToDnsStandardName(name);
        }
    }

    public class DnsPtrData
    {
        public string name { set; get; }
        public static DnsPtrData fromData(byte[] dataHead, byte[] data)
        {
            DnsPtrData obj = new DnsPtrData();
            obj.name = DnsMessage.parseDnsStandardName(dataHead, data);
            return obj;
        }
        public byte[] toData()
        {
            return DnsMessage.stringToDnsStandardName(name);
        }
    }

    public class DnsMessage
    {
        public DnsHeader header { get; set; }
        public List<DnsQuestion> questions { get; set; }
        public List<DnsResource> answers { get; set; }
        public List<DnsResource> nameServers { get; set; }
        public List<DnsResource> additionals { get; set; }
        class StdNameGenerator
        {
            class StdNameInfo
            {
                public string name;
                public long offset;
            }
            List<StdNameInfo> names = new List<StdNameInfo>();
            public byte[] generateName(string name)
            {
                MemoryStream s = new MemoryStream();
                BinaryWriter w = new BinaryWriter(s);
                while (name.Length > 0)
                {
                    foreach (StdNameInfo info in names)
                    {
                        var infoname = info.name;
                        if (infoname.EndsWith(name))
                        {
                            int index = infoname.Length - name.Length; // index in infoname.
                            if (index == 0 /*equal*/|| infoname[index - 1] == '.'/*subdomain*/)
                            {
                                long pos = info.offset + index;
                                w.Write((byte)(0xC0 | ((pos & 0xFF00) >> 8)));
                                w.Write((byte)(pos & 0xFF));
                                goto _return;
                            }
                        }
                    }
                    int n = name.IndexOf('.');
                    if (n != -1)
                    {
                        w.Write((byte)n);
                        w.Write(name.Substring(0, n).ToCharArray());
                        name = name.Substring(n + 1);
                    }
                    else
                    {
                        w.Write((byte)name.Length);
                        w.Write(name.ToCharArray());
                        name = "";
                    }
                }
                w.Write((byte)0);
                _return:
                return s.ToArray();
            }
            public void addName(string name, long offset)
            {
                var info = new StdNameInfo();
                info.name = name;
                info.offset = offset;
                this.names.Add(info);
            }
            public void generateAuto(string name, BinaryWriter w)
            {
                var bytes = generateName(name);
                addName(name, w.BaseStream.Position);
                w.Write(bytes);
            }
        }
        public DnsMessage()
        {
            header = new DnsHeader();
            questions = new List<DnsQuestion>();
            answers = new List<DnsResource>();
            nameServers = new List<DnsResource>();
            additionals = new List<DnsResource>();
        }
        public static ushort exchange2(ushort val)
        {
            return (ushort)(((val & 0xFF00) >> 8) | ((val & 0x00FF) << 8));
        }
        public static uint exchange4(uint val)
        {
            return (uint)(
                ((val & 0x000000FF) << 24) |
                ((val & 0x0000FF00) << 8) |
                ((val & 0x00FF0000) >> 8) |
                ((val & 0xFF000000) >> 24)
                );
        }
        //public static DnsData parse(byte[] data);
        public byte[] toData()
        {
            StdNameGenerator nameGenerator = new StdNameGenerator();
            MemoryStream s = new MemoryStream();
            BinaryWriter w = new BinaryWriter(s);
            writeHeader(w);

            foreach (DnsQuestion q in questions)
            {
                writeQuestion(q, w, nameGenerator);
            }

            foreach (DnsResource r in answers)
            {
                writeResource(r, w, nameGenerator);
            }
            foreach (DnsResource r in nameServers)
            {
                writeResource(r, w, nameGenerator);
            }
            foreach (DnsResource r in additionals)
            {
                writeResource(r, w, nameGenerator);
            }
            return s.ToArray();
        }
        private static void writeQuestion(DnsQuestion q, BinaryWriter w, StdNameGenerator nameGenerator)
        {
            //w.Write(stringToDnsStandardName(q.name));
            nameGenerator.generateAuto(q.name, w);
            w.Write(exchange2((ushort)q.recordType));
            w.Write(exchange2((ushort)q.dnsClass));
        }
        private static void writeResource(DnsResource r, BinaryWriter w, StdNameGenerator nameGenerator)
        {
            //w.Write(stringToDnsStandardName(r.name));
            nameGenerator.generateAuto(r.name, w);
            w.Write(exchange2((ushort)r.recordType));
            w.Write(exchange2((ushort)r.dnsClass));
            w.Write(exchange4(r.timeToLive));
            w.Write(exchange2(r.dataLength));
            w.Write(r.data);
        }
        private void writeHeader(BinaryWriter w)
        {
            w.Write(exchange2(header.id));
            w.Write(header.operationByte);
            w.Write(header.responseByte);
            w.Write(exchange2((ushort)questions.Count));
            w.Write(exchange2((ushort)answers.Count));
            w.Write(exchange2((ushort)nameServers.Count));
            w.Write(exchange2((ushort)additionals.Count));
        }
        public byte[] toUdpData()
        {
            byte[] data = new byte[512];
            toData().CopyTo(data, 0);
            return data;
        }

        internal static DnsMessage parse(byte[] data)
        {
            MemoryStream s = new MemoryStream(data);
            BinaryReader r = new BinaryReader(s);

            DnsMessage dnsData = new DnsMessage();
            DnsHeader header = parseHeader(r);
            dnsData.header = header;

            int questionCount = exchange2(r.ReadUInt16());
            int answerCount = exchange2(r.ReadUInt16());
            int nameServerCount = exchange2(r.ReadUInt16());
            int additionalCount = exchange2(r.ReadUInt16());

            while (questionCount-- != 0)
            {
                DnsQuestion question = parseQuestion(r, data);
                dnsData.questions.Add(question);
            }

            while (answerCount-- != 0)
            {
                DnsResource resource = parseResource(r, data);
                dnsData.answers.Add(resource);
            }

            while (nameServerCount-- != 0)
            {
                DnsResource resource = parseResource(r, data);
                dnsData.nameServers.Add(resource);
            }

            while (additionalCount-- != 0)
            {
                DnsResource resource = parseResource(r, data);
                dnsData.additionals.Add(resource);
            }

            return dnsData;
        }

        private static DnsHeader parseHeader(BinaryReader r)
        {
            DnsHeader header = new DnsHeader();
            header.id = exchange2(r.ReadUInt16());
            header.operationByte = r.ReadByte();
            header.responseByte = r.ReadByte();
            return header;
        }

        private static DnsQuestion parseQuestion(BinaryReader r, byte[] data)
        {
            DnsQuestion question = new DnsQuestion();
            question.name = parseDnsStandardName(r, data);
            question.recordType = (DnsRecordType)exchange2(r.ReadUInt16());
            question.dnsClass = (DnsClass)exchange2(r.ReadUInt16());
            return question;
        }

        private static DnsResource parseResource(BinaryReader r, byte[] data)
        {
            DnsResource resource = new DnsResource();
            resource.name = parseDnsStandardName(r, data);
            resource.recordType = (DnsRecordType)exchange2(r.ReadUInt16());
            resource.dnsClass = (DnsClass)exchange2(r.ReadUInt16());
            resource.timeToLive = exchange4(r.ReadUInt32());
            resource.dataLength = exchange2(r.ReadUInt16());
            switch (resource.recordType)
            {
                case DnsRecordType.CNAME:
                    resource.data = DnsCNameData.fromData(data, r.ReadBytes(resource.dataLength)).toData();
                    resource.dataLength = (ushort)resource.data.Length;
                    break;
                case DnsRecordType.NS:
                    resource.data = DnsNsData.fromData(data, r.ReadBytes(resource.dataLength)).toData();
                    resource.dataLength = (ushort)resource.data.Length;
                    break;
                default:
                    resource.data = r.ReadBytes(resource.dataLength);;
                    break;
            }
            return resource;
        }

        public static char[] separator = new char[] { '.', '@' };
        public static byte[] stringToDnsStandardName(string name)
        {
            var pieces = name.Split(separator);

            MemoryStream s = new MemoryStream();
            BinaryWriter w = new BinaryWriter(s);
            foreach (string str in pieces)
            {
                w.Write((byte)str.Length);
                w.Write(str.ToCharArray());
            }
            w.Write((byte)0);

            return s.ToArray();
        }
        public static string parseDnsStandardName(byte[] dataHead, byte[] data)
        {
            return parseDnsStandardName(new BinaryReader(new MemoryStream(data, 0, data.Length)), dataHead);
        }
        public static string parseDnsStandardName(System.IO.BinaryReader reader, byte[] dataHead)
        {
            string str = "";
            while (-1 != reader.PeekChar())
            {
                int length = reader.ReadByte();
                if ((length & 0xC0) == 0xC0)
                {
                    int low = reader.ReadByte();
                    int offset = ((length & ~0xC0) << 8) | low;
                    return str + parseDnsStandardName(new BinaryReader(new MemoryStream(dataHead, offset, dataHead.Length - offset)), dataHead);
                }
                else if (length > 0)
                {
                    while (length-- != 0)
                    {
                        str += Convert.ToChar(reader.ReadByte());
                    }
                    str += '.';
                }
                else /*if(length == 0)*/
                {
                    break;
                }
            }
            return str != "" ? str.Remove(str.Length - 1) : str;
        }

        //! deep copy
        public DnsMessage Clone()
        {
            var obj = new DnsMessage();
            obj.header = this.header.Clone();
            foreach (DnsQuestion q in this.questions)
            {
                obj.questions.Add(q.Clone());
            }

            foreach(DnsResource r in this.answers)
            {
                obj.answers.Add(r.Clone());
            }

            foreach (DnsResource r in this.nameServers)
            {
                obj.nameServers.Add(r.Clone());
            }

            foreach (DnsResource r in this.additionals)
            {
                obj.additionals.Add(r.Clone());
            }
            return obj;
        }

        public static DnsMessage queryDnsByUdp(DnsMessage msg, string server)
        {
            var ep = new System.Net.IPEndPoint(System.Net.IPAddress.Parse(server), 53);
            var u = new System.Net.Sockets.UdpClient(server, 53);
            var dataSent = msg.toUdpData();
            u.Send(dataSent, dataSent.Length);
            var dataReceived = u.Receive(ref ep);
            return DnsMessage.parse(dataReceived);
        }
    }
}
